package co.recloud.ariadne.thread;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentLinkedQueue;

import co.recloud.ariadne.model.Host;
import co.recloud.ariadne.model.Location;
import co.recloud.ariadne.model.logical.Column;
import co.recloud.ariadne.model.logical.Table;
import co.recloud.ariadne.model.logical.Transaction;
import co.recloud.ariadne.persistor.KeyKeyValuePersistor;
import co.recloud.ariadne.persistor.KeyKeyValuePersistorImpl;
import co.recloud.ariadne.persistor.KeyValuePersistor;
import co.recloud.ariadne.persistor.KeyValuePersistorImpl;
import co.recloud.ariadne.persistor.SchemaPersistor;
import co.recloud.ariadne.persistor.SchemaPersistorImpl;
import co.recloud.ariadne.request.KeyMigrateRequest;
import co.recloud.ariadne.response.Response;
import co.recloud.ariadne.store.AddressTable;
import co.recloud.ariadne.store.HostTable;
import co.recloud.ariadne.store.MemoryStore;

public class PartitionerThread extends Thread {

	public void run() {
		while (true) {
			if(!Main.waitForServiceStatus()) {
				continue;
			}
			if(!Main.waitForIdle()) {
				continue;
			}
			HostTable ht = HostTable.getInstance();
			MemoryStore store = MemoryStore.getStore(ht.getLocalhost()
					.getLocation());
			Set<String> paths = store.getAllPaths();
			Map<String, Map<String, Table>> columnToParentTable = new HashMap<String, Map<String, Table>>();
			Map<String, Set<Column>> tableToChildColumns = new HashMap<String, Set<Column>>();
			inferSchema(paths, columnToParentTable, tableToChildColumns);
			procesRelatedPaths(paths, columnToParentTable, tableToChildColumns);
		}
	}
	
	private void procesRelatedPaths(Set<String> paths, Map<String, Map<String, Table>> columnToParentTable, Map<String, Set<Column>> tableToChildColumns) {
		Set<String> migratedPaths = new HashSet<String>();
		Transaction transaction = new Transaction();
		transaction.generate();
		transaction.setStartTime(Main.getTime());
		for(String path : paths) {
			String schema;
			String columnFamily;
			String key;
			String[] keyTokens = path.split(":", 2);
			String[] schemaTokens = keyTokens[0].split("\\.", 2);
			key = keyTokens[1];
			columnFamily = schemaTokens[1];
			schema = schemaTokens[0];
			Set<String> relatedPaths = new HashSet<String>();
			if(schema.equals("key_values")) {
				KeyValuePersistor kv = new KeyValuePersistorImpl(
						transaction);
				KeyKeyValuePersistor kkv = new KeyKeyValuePersistorImpl(transaction);
				Table currentTable = new Table();
				currentTable.setName(columnFamily);
				Map<String, Object> row = kv.getAllColumns(key, currentTable);
				if(row != null) {
					for(String column : row.keySet()) {
						if(columnToParentTable.containsKey(columnFamily) && columnToParentTable.get(columnFamily).containsKey(column)) {
							Object value = row.get(column);
							if(value != null) {
								Table parentTable = columnToParentTable.get(columnFamily).get(column);
								if(parentTable != null) {
									String relatedPath = schema + "." + parentTable.getName() + ":" + value;
									relatedPaths.add(relatedPath);
								}
							}
						}
					}
					if(tableToChildColumns.containsKey(columnFamily) && tableToChildColumns.get(columnFamily) != null) {
						for(Column childColumn : tableToChildColumns.get(columnFamily)) {
							Set<String> childKeys = kkv.getParentToChild(new Column(currentTable, Table.PRIMARY_KEY), key, childColumn);
							for(String childKey : childKeys) {
								String relatedPath = schema + "." + childColumn.getTable().getName() + ":" + childKey;
								relatedPaths.add(relatedPath);
							}
						}
					}
				}
				processPath(path, relatedPaths, migratedPaths, transaction);
			}
		}
		transaction.setCommitTime(Main.getTime());
		transaction.commit();
	}

	private void processPath(String path, Set<String> relatedPaths, Set<String> migratedPaths, Transaction transaction) {
		Map<Long, Map<Long, Long>> edgesWithLocation = new HashMap<Long, Map<Long,Long>>();
		SortedMap<Long, Location> scores = new TreeMap<Long, Location>();
		HostTable ht = HostTable.getInstance();
		AddressTable at = AddressTable.getInstance();
		Location myLocation = ht.getLocalhost().getLocation();
		Long myX = myLocation.getX();
		Long myY = myLocation.getY();
		for(String relatedPath : relatedPaths) {
			String schema;
			String columnFamily;
			String key;
			String[] keyTokens = relatedPath.split(":", 2);
			String[] schemaTokens = keyTokens[0].split("\\.", 2);
			key = keyTokens[1];
			columnFamily = schemaTokens[1];
			schema = schemaTokens[0];
			Host otherHost = at.resolve(transaction, schema, columnFamily, key);
			Location otherLocation = otherHost.getLocation();
			Long otherX = otherLocation.getX();
			Long otherY = otherLocation.getY();
			if(!edgesWithLocation.containsKey(otherX)) {
				edgesWithLocation.put(otherX, new HashMap<Long, Long>());
			}
			if(!edgesWithLocation.get(otherX).containsKey(otherY)) {
				edgesWithLocation.get(otherX).put(otherY, 0L);
			}
			edgesWithLocation.get(otherX).put(otherY, edgesWithLocation.get(otherX).get(otherY) + 1);
		}
		Long myEdgeWeights = 0L;
		if(edgesWithLocation.containsKey(myX) && edgesWithLocation.get(myX).containsKey(myY)) {
			myEdgeWeights = edgesWithLocation.get(myX).get(myY);
		}
		for(Long otherX : edgesWithLocation.keySet()) {
			for(Long otherY : edgesWithLocation.get(otherX).keySet()) {
				Long score = 0L;
				Long edgeWeights = edgesWithLocation.get(otherX).get(otherY);
				Location otherLocation = new Location();
				otherLocation.setX(otherX);
				otherLocation.setY(otherY);
				if (myLocation.equals(otherLocation)) {
					score = 0L;
				} else if(ht.isReplica(otherLocation)) {
					score = myEdgeWeights;
				} else {
					score = edgeWeights - myEdgeWeights;
				}
				scores.put(score, otherLocation);
			}
		}
		if(!scores.isEmpty()) {
			MemoryStore store = MemoryStore.getStore(myLocation);
			Location bestLocation = scores.get(scores.lastKey());
			Queue<Response> responses = new ConcurrentLinkedQueue<Response>();
			if(!bestLocation.equals(myLocation)) {
				Set<Host> hosts = new HashSet<Host>();
				Host host = ht.getHostAtLocation(bestLocation);
				hosts.add(host);
				hosts.add(ht.getBestCoHost(host));
				store.condensePath(path);
				SortedMap<Long, SortedMap<Long, Map<String, Object>>> versions = store.getAllVersions(path);
				for(Host target : hosts) {
					KeyMigrateRequest keyReq = new KeyMigrateRequest();
					keyReq.setPath(path);
					keyReq.setVersions(versions);
					keyReq.setPrimary(target);
					keyReq.setTarget(target);
					keyReq.sendBlocking(target);
					if(keyReq.getResponse() != null) {
						responses.add(keyReq.getResponse());
					}
				}
				if(responses.size() == 2) {
					migratedPaths.add(path);
					at.setHost(transaction, path, host);
				}
			}
		}
	}

	private void inferSchema(Set<String> paths, Map<String, Map<String, Table>> columnToParentTable, Map<String, Set<Column>> tableToChildColumns) {
		for (String path : paths) {
			String schema;
			String columnFamily;
			String key;
			String[] keyTokens = path.split(":", 2);
			String[] schemaTokens = keyTokens[0].split("\\.", 2);
			key = keyTokens[1];
			columnFamily = schemaTokens[1];
			schema = schemaTokens[0];
			if (schema.equals("key_values")) {
				Transaction transaction = new Transaction();
				transaction.generate();
				transaction.setStartTime(Main.getTime());
				KeyValuePersistor kv = new KeyValuePersistorImpl(
						transaction);
				SchemaPersistor schemaPersistor = new SchemaPersistorImpl(
						transaction);
				if (!columnToParentTable.containsKey(columnFamily)) {
					columnToParentTable.put(columnFamily,
							new HashMap<String, Table>());
				}
				Table currentTable = new Table();
				currentTable.setName(columnFamily);
				Map<String, Object> row = kv.getAllColumns(key, currentTable);
				if (row != null) {
					for (String column : row.keySet()) {
						Column currentColumn = new Column(currentTable,
								column);
						Table parentTable = null;
						if (!columnToParentTable.get(columnFamily)
								.containsKey(column)) {
							parentTable = schemaPersistor
									.getParentTable(currentColumn);
							if (parentTable != null) {
								columnToParentTable.get(columnFamily).put(
										column, parentTable);
							} else {
								columnToParentTable.get(columnFamily).put(
										column, null);
							}
						}
						if(parentTable != null) {
							if(!tableToChildColumns.containsKey(parentTable.getName())) {
								tableToChildColumns.put(parentTable.getName(), new HashSet<Column>());
							}
							tableToChildColumns.get(parentTable.getName()).add(currentColumn);
						}
					}
				}
				transaction.setCommitTime(Main.getTime());
				transaction.commit();
			}
		}
	}
}
